import os
import math
import random
import numpy as np
from enum import Enum
from PIL import Image, ImageDraw, ImageFont, ImageColor


class BackgroundGradient(Enum):
    skyline = ['#1488CC', '#2B32B2']
    cool_brown = ['#603813', '#b29f94']
    rose_water = ['#E55D87', '#5FC3E4']
    crystal_clear = ['#159957', '#155799']


class BackgroundShape(Enum):
    circle = 1
    triangle = 2
    square = 3
    hexagon = 4


class BlogCoverGenerator:

    def __init__(self,
                 title: str,
                 sub_title: str,
                 title_h_ratio: float=.5,
                 ratio_pair: tuple[int]=(16, 9),
                 bg_gradient: BackgroundGradient=BackgroundGradient.skyline,
                 bg_shape: BackgroundShape=BackgroundShape.circle,
                 min_size='1M'):
        # 封面主标题
        self.title = title
        # 封面副标题
        self.sub_title = sub_title
        # 如果没有副标题，主标题需要垂直居中
        self.no_sub_title = False
        if self.sub_title is None or '' == self.sub_title.strip():
            self.no_sub_title = True
        # 标题区域垂直方向占比(从正中心开始计算)，参考值0.3~0.6
        self.title_h_ratio = title_h_ratio
        if self.title_h_ratio < 0.3:
            self.title_h_ratio = 0.3
        elif self.title_h_ratio > 0.6:
            self.title_h_ratio = 0.6
        # 字体位置
        self.zh_font_location = ''
        self.en_font_location = ''
        # 封面宽高比，16:9, 4:3, 1:1等，以16:9为例，需要传入(16, 9)
        self.ratio_pair = ratio_pair
        # 封面渐变背景色
        self.bg_gradient = bg_gradient
        # 封面背景几何图形，目前支持三角形、六边形、圆形
        self.bg_shape = bg_shape
        # 封面大小。如果传入字符串，支持的单位为k, M；也可以传入数值
        self.min_size = self._parse_min_size(min_size)
        if self.min_size > 178956970:
            # ImageDraw.text大小限制
            self.min_size = 178956970

    @staticmethod
    def _parse_min_size(min_size: object):
        if type(min_size) == int:
            return min_size
        pre_min_size_str = min_size.strip().lower()
        if pre_min_size_str.endswith('k'):
            num_str = pre_min_size_str.split('k')[0]
            return 1024 * int(num_str)
        elif pre_min_size_str.endswith('m'):
            num_str = pre_min_size_str.split('m')[0]
            return 1024 ** 2 * int(num_str)
        # 不带单位，直接返回数值
        return int(pre_min_size_str)

    def _get_cover_size(self):
        # 先按比例求出height
        height_square = self.min_size * self.ratio_pair[1] / self.ratio_pair[0]
        height = int(math.sqrt(height_square)) + 1

        # 适当增大height，让height和width严格按照ratio的比例，并且都为整数
        for i in range(self.min_size):
            if (height + i) % self.ratio_pair[1] == 0:
                height += i
                break

        width = int(height * self.ratio_pair[0] / self.ratio_pair[1])
        return width, height

    def _display_gradient_bg(self, base_img: Image):
        img_w, img_h = base_img.size
        draw = ImageDraw.Draw(base_img)
        start_color, end_color = self.bg_gradient.value
        if '#' in start_color:
            start_color = ImageColor.getrgb(start_color)
        if '#' in end_color:
            end_color = ImageColor.getrgb(end_color)

        # 水平方向渐变，渐变步长
        step_r = (end_color[0] - start_color[0]) / img_w
        step_g = (end_color[1] - start_color[1]) / img_w
        step_b = (end_color[2] - start_color[2]) / img_w

        for i in range(0, img_w):
            bg_r = round(start_color[0] + step_r * i)
            bg_g = round(start_color[1] + step_g * i)
            bg_b = round(start_color[2] + step_b * i)
            draw.line([(i, 0), (i, img_h)], fill=(bg_r, bg_g, bg_b))

    @staticmethod
    def _width_regular_polygon(draw: ImageDraw, width: int,
                               bounding_circle, n_sides, rotation=0, fill=None, outline=None):
        """
        pillow提供的regular_polygon，不支持对outline设置width，自定义方法，支持轮廓宽度
        """
        start = bounding_circle[2]
        for i in np.arange(0, width, 0.05):
            new_bounding_circle = (bounding_circle[0], bounding_circle[1], start + i)
            draw.regular_polygon(bounding_circle=new_bounding_circle, n_sides=n_sides,
                                 rotation=rotation, fill=fill, outline=outline)

    def _display_decorate_shape(self, base_img: Image):
        """
        封面边缘填充装饰图形，右上角一个大图形，1个小图形，左下角3个小图形
        :param base_img: 封面图片
        """
        def draw_one_circle(the_draw: ImageDraw, bounding_circle, shape_color, width):
            the_draw.ellipse([
                (bounding_circle[0] - bounding_circle[2], bounding_circle[1] - bounding_circle[2]),
                (bounding_circle[0] + bounding_circle[2], bounding_circle[1] + bounding_circle[2])
            ], outline=shape_color, width=width)

        def draw_all_regular_polygon(width_regular_polygon, n_size: int):
            # 右上角
            rotation_1, rotation_2 = 15, 30
            if n_size == 3:
                rotation_2 = 180
            elif n_size == 4:
                rotation_1 = 30
            elif n_size == 6:
                rotation_1 = 0
                rotation_2 = 30
            width_regular_polygon(draw, 9, bounding_circle=b_circle_1, n_sides=n_size,
                                  rotation=rotation_1, outline=shape_color_1)
            width_regular_polygon(draw, 3, bounding_circle=b_circle_2, n_sides=n_size,
                                  rotation=rotation_2, outline=shape_color_2)
            # 左下角
            width_regular_polygon(draw, 12, bounding_circle=b_circle_3, n_sides=n_size,
                                  rotation=random.randint(30, 75), outline=shape_color_2)
            width_regular_polygon(draw, 3, bounding_circle=b_circle_4, n_sides=n_size,
                                  rotation=random.randint(40, 80), outline=shape_color_3)
            width_regular_polygon(draw, 6, bounding_circle=b_circle_5, n_sides=n_size,
                                  rotation=random.randint(90, 120), outline=shape_color_2)
            width_regular_polygon(draw, 3, bounding_circle=b_circle_6, n_sides=n_size,
                                  rotation=random.randint(10, 60), outline=shape_color_3)

        img_w, img_h = base_img.size
        img_shape = Image.new('RGBA', (img_w, img_h), color=(0, 0, 0, 0))
        draw = ImageDraw.Draw(img_shape)
        max_size = max(img_w * (1 - self.title_h_ratio), img_h * (1 - self.title_h_ratio))
        shape_color_1 = (255, 255, 255, 120)
        shape_color_2 = (255, 255, 255, 60)
        shape_color_3 = (255, 255, 255, 45)
        # 右上角2个图形
        r_1 = int(max_size / 2)
        b_circle_1 = (img_w, 0, r_1)
        r_2 = r_1
        delta_2 = int(r_1 / 3)
        b_circle_2 = (img_w - r_2 - delta_2, delta_2 - r_2, r_2)

        # 左下角4个图形
        r_3 = int(max_size / 12)
        b_circle_3 = (r_3 * 2, img_h - 2 * r_3, r_3)
        r_4 = int(max_size / 6)
        b_circle_4 = (r_3 + 1.2 * r_4, img_h - 2 * r_4, r_4)
        r_5 = int(max_size / 18)
        delta_5 = random.randint(3, r_5)
        b_circle_5 = (r_5 + delta_5, img_h - 9 * r_5, r_5)
        r_6 = int(max_size / 20)
        delta_6 = random.randint(6, 3 * r_6)
        b_circle_6 = (13 * r_6, img_h - r_6 - delta_6, r_6)

        if BackgroundShape.circle == self.bg_shape:
            # 右上角
            draw_one_circle(draw, b_circle_1, shape_color_1, 9)
            draw_one_circle(draw, b_circle_2, shape_color_3, 3)

            # 左下角
            draw_one_circle(draw, b_circle_3, shape_color_2, 12)
            draw_one_circle(draw, b_circle_4, shape_color_3, 3)
            draw_one_circle(draw, b_circle_5, shape_color_2, 6)
            draw_one_circle(draw, b_circle_6, shape_color_3, 3)
        elif BackgroundShape.triangle == self.bg_shape:
            draw_all_regular_polygon(self._width_regular_polygon, 3)
        elif BackgroundShape.square == self.bg_shape:
            draw_all_regular_polygon(self._width_regular_polygon, 4)
        elif BackgroundShape.hexagon == self.bg_shape:
            draw_all_regular_polygon(self._width_regular_polygon, 6)

        return Image.alpha_composite(base_img, img_shape)

    @staticmethod
    def _display_transparent_mask(base_img: Image):
        img_w, img_h = base_img.size
        img_mask = Image.new('RGBA', (img_w, img_h), color=(0, 0, 0, 135))
        return Image.alpha_composite(base_img, img_mask)

    def _get_real_font(self, target_title: str, font_size: int):
        for ch in target_title:
            if u'\u4e00' <= ch <= u'\u9fff':
                # 中文字体
                return ImageFont.truetype(self.zh_font_location, font_size)
        # 英文字体
        return ImageFont.truetype(self.en_font_location, font_size)

    def _display_title(self, base_img: Image):
        def get_checked_font_size(target_title: str, font_size: int, max_width: int):
            # 预检查，判断文字宽度是否超出封面
            check_font = self._get_real_font(target_title, font_size)
            _, _, check_w, check_h = check_font.getbbox(target_title)
            if check_w > max_width:
                scale_ratio = max_width / check_w
                font_size = int(font_size * scale_ratio)
            return font_size

        def get_best_font_by_title(target_title: str, max_width: int, max_height: int, is_sub: bool=False):
            font_size = max_height
            if self.no_sub_title:
                font_size = get_checked_font_size(target_title, font_size, max_width)
            else:
                spacing = 12
                # 两个标题都存在，分配比例3(主) : 1(副)，间隔一定距离
                if is_sub:
                    # 副标题
                    font_size = int((max_height - spacing) / 4)
                    font_size = get_checked_font_size(target_title, font_size, max_width)
                else:
                    # 主标题
                    font_size = int(3 * (max_height - spacing) / 4)
                    font_size = get_checked_font_size(target_title, font_size, max_width)
            return self._get_real_font(target_title, font_size)

        title_d = ImageDraw.Draw(base_img)

        img_w, img_h = base_img.size
        text_h = int(img_h * self.title_h_ratio)
        text_h_up = (img_h - text_h) / 2
        text_h_down = text_h_up + text_h
        text_max_w = int(img_w * 0.8)

        font_title = get_best_font_by_title(self.title, text_max_w, text_h)
        if self.no_sub_title:
            # 绘制主标题
            title_d.text((img_w / 2, img_h / 2), self.title, font=font_title, fill='white', anchor='mm')
        else:
            # 确定文字高度
            _, _, font_title_w, font_title_h = font_title.getbbox(self.title)
            font_sub_title = get_best_font_by_title(self.sub_title,
                                                    min(font_title_w, text_max_w),
                                                    text_h, is_sub=True)
            _, _, font_sub_title_w, font_sub_title_h = font_sub_title.getbbox(self.sub_title)
            title_spacing = 18
            optimized_text_h = font_title_h + font_sub_title_h + title_spacing
            if optimized_text_h < text_h:
                text_h_up = (img_h - optimized_text_h) / 2
                text_h_down = text_h_up + optimized_text_h

            # 绘制主标题
            title_d.text((img_w / 2, text_h_up), self.title, font=font_title, fill='white', anchor='mt')

            # 绘制副标题，宽度不能超过主标题
            title_d.text((img_w / 2, text_h_down), self.sub_title, font=font_sub_title,
                         fill='white', anchor='mb')

            # 主标题和副标题之间绘制一条透明细线
            line_img = Image.new('RGBA', (img_w, img_h), color=(0, 0, 0, 0))
            draw = ImageDraw.Draw(line_img)
            line_y_pos = text_h_up + font_title_h
            line_x_pos = (img_w - font_title_w) / 2
            draw.line([(line_x_pos, line_y_pos), (img_w - line_x_pos, line_y_pos)],
                      fill=(255, 255, 255, 120), width=2)
            base_img = Image.alpha_composite(base_img, line_img)
        return base_img

    def set_fonts(self, zh_font_location: str, en_font_location: str):
        self.zh_font_location = zh_font_location
        self.en_font_location = en_font_location

    def generate_cover(self, output_path: str, output_file_name: str):
        width, height = self._get_cover_size()
        img = Image.new('RGBA', (width, height))

        # 第1层，渐变背景
        self._display_gradient_bg(img)

        # 第2层，边缘装饰图形
        img = self._display_decorate_shape(img)

        # 第3层，半透明遮罩
        img = self._display_transparent_mask(img)

        # 第4层，文字
        img = self._display_title(img)

        # 保存
        img.save(os.path.join(output_path, output_file_name))

def generator_image(cover_title,cover_sub_title,movie_name):
    cover_title = '三国演义'
    #cover_sub_title = '第一部 群雄逐鹿 桃园三结义'
    generator = BlogCoverGenerator(cover_title, cover_sub_title, title_h_ratio=0.5,
                                   bg_gradient=BackgroundGradient.rose_water,
                                   bg_shape=BackgroundShape.circle,
                                   ratio_pair=(4, 3), min_size='3M')
    generator.set_fonts('fonts/HYWenHei-85W-1.ttf', 'fonts/Kannada MN.ttc')
    generator.generate_cover('./cover/'+movie_name, cover_sub_title+'.png')
if __name__ == '__main__':
    cover_title = '三国演义'
    cover_sub_title = '第一部 群雄逐鹿 桃园三结义'
    generator = BlogCoverGenerator(cover_title, cover_sub_title, title_h_ratio=0.5,
                                   bg_gradient=BackgroundGradient.rose_water,
                                   bg_shape=BackgroundShape.circle,
                                   ratio_pair=(4, 3), min_size='3M')
    generator.set_fonts('fonts/HYWenHei-85W-1.ttf', 'fonts/Kannada MN.ttc')
    generator.generate_cover('./cover', cover_title+'_'+cover_sub_title+'.png')

